#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

#include <cuda.h>
#include <cuda_runtime.h>

#include <vector>

#define STATUS_WIN 0
#define STATUS_OPEN4 1
#define STATUS_DEAD4 2
#define STATUS_OPEN3 3
#define STATUS_SLEEP3 4
#define STATUS_OPEN2 5

namespace {
__global__ void check_status_cuda_kernel(
    torch::PackedTensorAccessor<bool,4,torch::RestrictPtrTraits,size_t> board,
    torch::PackedTensorAccessor<bool,4,torch::RestrictPtrTraits,size_t> status){
    const static int dirs[4][2] = {{1, 0}, {0, 1}, {1, 1}, {1, -1}};
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;

    auto h = board.size(2);
    auto w = board.size(3);
    auto i = idx / (w * h);
    auto j = (idx / w) % h;
    auto k = idx % w;
    
    if (i >= board.size(0)) return;
    if (board[i][0][j][k]) return;
    if (board[i][1][j][k]) return;

    for (int side = 0; side < 2; side++) {
      for (auto& dir : dirs) {
        int count = 1;
        int open_cnt = 0;
        int x = j + dir[0], y = k + dir[1];
        while (x >= 0 && y >= 0 && x < board.size(2) && y < board.size(3) &&
              board[i][side][x][y])
          x += dir[0], y += dir[1], count++;
        if (x >= 0 && y >= 0 && x < board.size(2) && y < board.size(3) &&
              ~board[i][1-side][x][y])
          open_cnt += 1;
        x = j - dir[0], y = k - dir[1];
        while (x >= 0 && y >= 0 && x < board.size(2) && y < board.size(3) &&
              board[i][side][x][y])
          x -= dir[0], y -= dir[1], count++;
        if (x >= 0 && y >= 0 && x < board.size(2) && y < board.size(3) &&
              ~board[i][1-side][x][y])
          open_cnt += 1;
        if (count >= 5) {
          status[i][side*6+STATUS_WIN][j][k] = 1;
        }
        else if (count == 4) {
          if (open_cnt == 2) {
            status[i][side*6+STATUS_OPEN4][j][k] = 1;
          }
          if (open_cnt == 1) {
            status[i][side*6+STATUS_DEAD4][j][k] = 1;
          }
        }
        else if (count == 3) {
          if (open_cnt == 2) {
            status[i][side*6+STATUS_OPEN3][j][k] = 1;
          }
          if (open_cnt == 1) {
            status[i][side*6+STATUS_SLEEP3][j][k] = 1;
          }
        } else if (count == 2) {
          if (open_cnt == 2) {
            status[i][side*6+STATUS_OPEN2][j][k] = 1;
          }
        }
      }
    }
}
}

torch::Tensor get_status(torch::Tensor board){
    const int total_threads = board.size(0) * board.size(2) * board.size(3);
    const int threads = min(1024, total_threads);
    const dim3 blocks((total_threads + threads - 1) / threads);

    torch::Tensor status = torch::zeros(
        {board.size(0), 12, board.size(2), board.size(3)}, 
        torch::TensorOptions().dtype(torch::kBool).device(board.device())
    );

    auto stream = at::cuda::getCurrentCUDAStream();
    check_status_cuda_kernel<<<blocks, threads, 0, stream>>>(
        board.packed_accessor<bool,4,torch::RestrictPtrTraits,size_t>(),
        status.packed_accessor<bool,4,torch::RestrictPtrTraits,size_t>());
    return status;
}
